import time
import requests
from DrissionPage import ChromiumPage, ChromiumOptions, WebPage, Chromium
from urllib.parse import urlparse, parse_qs
import json
import os
import re
import tkinter as tk
from tkinter import ttk, scrolledtext, messagebox
import threading


# -------------------- 浏览器连接处理 --------------------
def initialize_browser_connection():
    browser = None
    try:
        # 创建 Chromium 配置对象
        # 指定浏览器可执行文件路径（项目根目录下的chrome/chrome.exe）
        # co = ChromiumOptions()
        # chrome_path = os.path.join(os.getcwd(), 'chrome', 'chrome.exe')
        # co.set_browser_path(chrome_path)
        # 创建浏览器连接
        browser = Chromium()
        tab = browser.latest_tab

        # 访问网页
        tab.get('https://mp.weixin.qq.com/')

        # 等待页面加载或URL变化
        start_time = time.time()
        timeout = 30  # 30秒超时

        while time.time() - start_time < timeout:
            current_url = tab.url
            if 'mp.weixin.qq.com' not in current_url:
                print("检测到URL已变化，关闭浏览器")
                break

            # 检查是否已登录（可以根据实际需求调整条件）
            if 'token=' in current_url:
                print("获取到所需的信息，请保持当前状态")
                return browser  # 返回浏览器实例供后续使用

            time.sleep(1)  # 每秒检查一次

        print("浏览器初始化完成或超时")
        return browser

    except Exception as e:
        print(f"浏览器初始化失败: {e}")
        if browser:
            browser.quit()
        return None


# -------------------- 原有核心函数 --------------------
def get_token():
    co = (
        ChromiumOptions()
        .headless()
        .no_imgs()
        .set_argument("--no-sandbox")
        .set_argument("--single-process")
    )
    page = WebPage(chromium_options=co)
    try:
        page.get('https://mp.weixin.qq.com/')
        final_url = page.url
        parsed = urlparse(final_url)
        token = parse_qs(parsed.query).get('token', [None])[0]
        return token
    except Exception as e:
        print(f"发生错误: {e}")
        return None
    finally:
        page.quit()


def get_cookies():
    co = ChromiumOptions()
    co.headless()
    page = ChromiumPage(co)
    page.get('https://mp.weixin.qq.com/')
    cookies = json.dumps(page.cookies())
    data = json.loads(cookies)
    result = {}
    for item in data:
        result[item["name"]] = item["value"]
    return result


def write_txt(content_url):
    co = ChromiumOptions()
    co.headless().no_imgs(True).mute(True)
    browser = Chromium(co)
    tab = browser.latest_tab
    tab.get(content_url)
    try:
        title = tab.ele("x://div[@id='img-content']/h1[@id='activity-name']").text
        date_time = tab.ele("x://span[@id='meta_content_hide_info']/em[@id='publish_time']").text
        ip_location = tab.ele("x://span[@id='js_ip_wording']").text
        main_content = "".join(tab.ele("x://div[@id='img-content']/div[@id='js_content']").text).strip().replace('\n',
                                                                                                                 '')
        gzh_name = tab.ele("x://a[@id='js_name']").text
        file_name = safe_filename(title)
        # 如果目录不存在则进行创建目录

        if not os.path.exists(f"{gzh_name}/{file_name}"):
            os.makedirs(f"{gzh_name}/{file_name}")
        with open(f"{gzh_name}/{file_name}/{file_name}.txt", "w", encoding="utf-8") as f:
            f.write(title + "\n\n")
            f.write(date_time + "   ")
            f.write(ip_location + "\n\n")
            f.write(main_content)
        # 对整页截图并保存
        tab.get_screenshot(path=f'{gzh_name}/{file_name}', name=f'{file_name}.jpg', full_page=False)
    except:
        pass


def safe_filename(title):
    safe_name = re.sub(r'[\\/*?:"<>| ]', '_', title)
    safe_name = re.sub(r'_+', '_', safe_name).strip('_')
    return safe_name if safe_name else 'unnamed_file'


# -------------------- GUI 界面实现 --------------------
class WeChatDownloader:
    def __init__(self, master):
        self.master = master
        master.title("微信公众号文章下载器")
        master.geometry("800x600")

        # 初始化浏览器参数
        self.token = None
        self.cookies = None
        self.headers = {
            "user-agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/134.0.0.0 Safari/537.36"
        }

        # UI组件
        self.create_widgets()
        self.initialize_browser()

    def create_widgets(self):
        # 输入区域
        input_frame = ttk.Frame(self.master)
        input_frame.pack(pady=10, fill=tk.X)

        self.entry = ttk.Entry(input_frame, width=40)
        self.entry.pack(side=tk.LEFT, padx=5)
        self.entry.bind("<Return>", self.start_download)

        self.btn = ttk.Button(input_frame, text="开始下载", command=self.start_download)
        self.btn.pack(side=tk.LEFT, padx=5)

        # 日志区域
        self.log_area = scrolledtext.ScrolledText(self.master, wrap=tk.WORD)
        self.log_area.pack(pady=10, padx=10, fill=tk.BOTH, expand=True)
        self.log_area.configure(state='disabled')

        # 状态栏
        self.status_var = tk.StringVar()
        ttk.Label(self.master, textvariable=self.status_var).pack(side=tk.BOTTOM, fill=tk.X)

    def initialize_browser(self):
        def init_task():
            self.log("正在初始化浏览器组件,请稍后...")
            self.token = get_token()
            self.cookies = get_cookies()
            self.log("准备就绪，请输入公众号名称")
            self.status_var.set("就绪")

        threading.Thread(target=init_task, daemon=True).start()

    def log(self, message):
        self.log_area.configure(state='normal')
        self.log_area.insert(tk.END, message + "\n")
        self.log_area.see(tk.END)
        self.log_area.configure(state='disabled')
        self.master.update()

    def start_download(self, event=None):
        query = self.entry.get().strip()
        if not query:
            messagebox.showwarning("输入错误", "请输入公众号名称")
            return

        self.btn.config(state='disabled')
        self.entry.config(state='disabled')
        self.status_var.set("正在下载...")

        def download_task():
            try:
                fakeid = self.search_key(query)
                if not fakeid:
                    self.log("未找到该公众号")
                    return

                self.log(f"\n开始下载公众号：{query}")
                page = 1
                while True:
                    params = {
                        "sub": "list",
                        'begin': (page - 1) * 5,
                        "count": "5",
                        "fakeid": fakeid,
                        "token": self.token,
                        "f": "json",
                    }
                    response = requests.get(
                        "https://mp.weixin.qq.com/cgi-bin/appmsgpublish",
                        headers=self.headers,
                        cookies=self.cookies,
                        params=params
                    )

                    if response.status_code != 200:
                        self.log(f"第 {page} 页请求失败，状态码: {response.status_code}")
                        break

                    try:
                        data = response.json()
                        page_data = json.loads(data['publish_page'])
                        if not page_data["publish_list"]:
                            self.log("没有更多数据了，程序已结束---")
                            break

                        self.log(f"\n提取到的第{page}页的文章：")
                        num = 0
                        for item in page_data["publish_list"]:
                            publish_info = json.loads(item["publish_info"])
                            for msg in publish_info["appmsg_info"]:
                                self.log(f"正在下载：{msg['title']}")
                                write_txt(msg["content_url"])
                                num += 1
                                time.sleep(1)
                        self.log(f"第{page}页完成{num}篇下载")
                        page += 1
                    except Exception as e:
                        self.log(f"解析错误: {str(e)}")
                        break
                    time.sleep(2)
            except Exception as e:
                self.log(f"发生错误：{str(e)}")
            finally:
                self.btn.config(state='normal')
                self.entry.config(state='normal')
                self.status_var.set("下载完成")

        threading.Thread(target=download_task, daemon=True).start()

    def search_key(self, query):
        url = "https://mp.weixin.qq.com/cgi-bin/searchbiz"
        params = {
            "action": "search_biz",
            "count": "1",
            "query": query,
            "token": self.token,
        }
        try:
            response = requests.get(url, headers=self.headers, cookies=self.cookies, params=params)
            return response.json()['list'][0]['fakeid']
        except:
            return None


if __name__ == '__main__':
    # 使用示例
    browser = initialize_browser_connection()
    if not browser:
        # 直接退出主程序
        print("浏览器初始化失败，请检查登录情况是否正常")
    root = tk.Tk()
    app = WeChatDownloader(root)
    root.mainloop()
